#!/usr/bin/python

from bcc import BPF
# from bcc.utils import printb
import sys
import socket
import struct
import time
import ctypes

def usage():
    print("Usage: {0} <ifdev> ".format(sys.argv[0]))
    print("e.g.: {0} eth0 (default eth0) \n".format(sys.argv[0]))

if len(sys.argv) != 2:
    usage()
    ifdev = "eth0"
else:
    ifdev = sys.argv[1]

# load BPF program
b = BPF(text = """
#include <uapi/linux/bpf.h>
#include <linux/in.h>
#include <linux/ip.h>
#include <linux/if_ether.h>

struct key_v4 {
    __u32 prefixlen;
    __u32 saddr;
};

BPF_LPM_TRIE(blacklist, struct key_v4, u32, 16);

        
int firewall(struct xdp_md *ctx) {
    uint64_t nh_off = 0;
    void *data_end = (void *)(long)ctx->data_end;
    void *data = (void *)(long)ctx->data;

    // Only IPv4 supported for this example
    struct ethhdr *ether = data;
    if (data + sizeof(*ether) > data_end) {
        // Malformed Ethernet header
        return XDP_ABORTED;
    }

    if (ether->h_proto != 0x08U) {  // htons(ETH_P_IP) -> 0x08U
        // Non IPv4 traffic
        return XDP_PASS;
    }

    data += sizeof(*ether);
    struct iphdr *ip = data;
    nh_off = sizeof(*ip);
    if (data + nh_off > data_end) {
        // Malformed IPv4 header
        return XDP_ABORTED;
    }


    // Lookup SRC IP in blacklisted IPs
    struct key_v4 key = {32, ip->saddr};
    __u64 *rule_idx = blacklist.lookup(&key);
    if (rule_idx) {
        // Matched, increase match counter for matched "rule"
        __u32 index = *(__u32*)rule_idx;  // make verifier happy
        bpf_trace_printk("XDP_DROP %d\\n", ip->saddr);
        return XDP_DROP;
    }

    // bpf_trace_printk("XDP_PASS %d\\n", ip->saddr);
    return XDP_PASS;
}
""", cflags=["-w"])


fn = b.load_func("firewall", BPF.XDP)
b.attach_xdp(ifdev, fn)

# 获取黑名单 map
blacklist = b["blacklist"]


# 定义一个 ctypes 结构体来存储 IP 地址和前缀长度
class LpmKey(ctypes.Structure):
    _fields_ = [("prefixlen", ctypes.c_uint32),
                ("addr", ctypes.c_uint32)]
    

# 添加 IP 地址到黑名单
ip = socket.inet_aton("172.31.73.80")  # 将 IP 地址转换为网络字节序
# print(struct.unpack("I", ip)[0])
key = LpmKey(32, struct.unpack("I", ip)[0])
blacklist[key] = ctypes.c_ubyte(1)

print("Blacklist loaded, hit CTRL+C to stop. Monitoring packets...")
while 1:
    try:
        # (task, pid, cpu, flags, ts, msg) = b.trace_fields()
        time.sleep(1)
        b.trace_print()
    except KeyboardInterrupt:
        print("Removing filter from device")
        break
    # printb(b"%-18.9f %-16s %-6d %s" % (ts, task, pid, msg))

b.remove_xdp(ifdev)

